import torch
from classifier.continual import ContinualClf
from .base import Base

torch.__version__


class BaseSingle(Base):
    fc: ContinualClf
    fc_func: str

    def __init__(
        self,
        args,
    ):
        super().__init__(
            args,
        )
        self.fc_func = args.get("fc_func", None)
        self.constant_dim = True

    def forward(self, x, ca=False):
        if ca:
            return self.fc(x)

        x = self.backbone(x)
        out = self.fc(x)
        out.update({"features": x})

        return out

    def extract_token(self, x):
        return self.backbone(x)

    def extract_feats(self, x):
        return self.backbone.forward_feats(x)  # type: ignore


SSIAT = BaseSingle
